# Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from typing import TYPE_CHECKING, Literal, Optional, Union

import numpy as np
from datasets import Dataset, load_dataset, load_from_disk

from ..extras import logging
from ..extras.constants import FILEEXT2TYPE
from ..extras.misc import check_version, has_tokenized_data
from .converter import align_dataset
from .data_utils import get_dataset_module, merge_dataset, read_cloud_json, split_dataset
from .parser import get_dataset_list
from .processor import (
    FeedbackDatasetProcessor,
    PackedSupervisedDatasetProcessor,
    PairwiseDatasetProcessor,
    PretrainDatasetProcessor,
    SupervisedDatasetProcessor,
    UnsupervisedDatasetProcessor,
)


if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset
    from transformers import PreTrainedTokenizer, ProcessorMixin, Seq2SeqTrainingArguments

    from ..hparams import DataArguments, ModelArguments
    from .data_utils import DatasetModule
    from .parser import DatasetAttr
    from .processor import DatasetProcessor
    from .template import Template


logger = logging.get_logger(__name__)


# 加载数据集
def _load_single_dataset(
    dataset_attr: "DatasetAttr",
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
) -> Union["Dataset", "IterableDataset"]:
    r"""Load a single dataset and aligns it to the standard format."""
    # 记录正在加载的数据集信息
    logger.info_rank0(f"Loading dataset {dataset_attr}...")
    # 初始化数据路径、名称、目录、文件列表变量
    data_path, data_name, data_dir, data_files = None, None, None, None
    # 如果是来自 HuggingFace / ModelScope / OpenMind 等模型库，设置对应参数
    if dataset_attr.load_from in ["hf_hub", "ms_hub", "om_hub"]:
        data_path = dataset_attr.dataset_name
        data_name = dataset_attr.subset
        data_dir = dataset_attr.folder

    # 如果是本地脚本方式加载，则拼接本地路径
    elif dataset_attr.load_from == "script":
        data_path = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)
        data_name = dataset_attr.subset
        data_dir = dataset_attr.folder

    # 如果是云端文件加载，直接使用指定 URL
    elif dataset_attr.load_from == "cloud_file":
        data_path = dataset_attr.dataset_name

    # 如果是本地文件加载，判断是否为目录或文件，并收集所有匹配文件
    elif dataset_attr.load_from == "file":
        data_files = []
        local_path = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)
        # 是目录则遍历所有文件
        if os.path.isdir(local_path):  # is directory
            for file_name in os.listdir(local_path):
                data_files.append(os.path.join(local_path, file_name))
        # 是文件则直接添加
        elif os.path.isfile(local_path):  # is file
            data_files.append(local_path)
        else:
            raise ValueError(f"File {local_path} not found.")

        # 获取第一个文件的扩展名并映射为 HF Dataset 支持的类型
        data_path = FILEEXT2TYPE.get(os.path.splitext(data_files[0])[-1][1:], None)
        if data_path is None:
            raise ValueError("Allowed file types: {}.".format(",".join(FILEEXT2TYPE.keys())))

        # 确保所有文件类型一致
        if any(data_path != FILEEXT2TYPE.get(os.path.splitext(data_file)[-1][1:], None) for data_file in data_files):
            raise ValueError("File types should be identical.")
    else:
        raise NotImplementedError(f"Unknown load type: {dataset_attr.load_from}.")

    # 如果是 ModelScope Hub 加载，使用 MsDataset 并转为 HuggingFace 格式
    if dataset_attr.load_from == "ms_hub":
        check_version("modelscope>=1.11.0", mandatory=True)
        from modelscope import MsDataset  # type: ignore
        from modelscope.utils.config_ds import MS_DATASETS_CACHE  # type: ignore

        cache_dir = model_args.cache_dir or MS_DATASETS_CACHE
        dataset = MsDataset.load(
            dataset_name=data_path,
            subset_name=data_name,
            data_dir=data_dir,
            data_files=data_files,
            split=dataset_attr.split,
            cache_dir=cache_dir,
            token=model_args.ms_hub_token,
            use_streaming=data_args.streaming,
        )
        if isinstance(dataset, MsDataset):
            dataset = dataset.to_hf_dataset()

    # 如果是 OpenMind Hub 加载，使用 OmDataset
    elif dataset_attr.load_from == "om_hub":
        check_version("openmind>=0.8.0", mandatory=True)
        from openmind import OmDataset  # type: ignore
        from openmind.utils.hub import OM_DATASETS_CACHE  # type: ignore

        cache_dir = model_args.cache_dir or OM_DATASETS_CACHE
        dataset = OmDataset.load_dataset(
            path=data_path,
            name=data_name,
            data_dir=data_dir,
            data_files=data_files,
            split=dataset_attr.split,
            cache_dir=cache_dir,
            token=model_args.om_hub_token,
            streaming=data_args.streaming,
        )
    # 如果是云端 JSON 文件，调用自定义函数读取
    elif dataset_attr.load_from == "cloud_file":
        dataset = Dataset.from_list(read_cloud_json(data_path), split=dataset_attr.split)
    else:
        # 默认使用 HuggingFace 的 load_dataset 方法加载数据
        dataset = load_dataset(
            path=data_path,
            name=data_name,
            data_dir=data_dir,
            data_files=data_files,
            split=dataset_attr.split,
            cache_dir=model_args.cache_dir,
            token=model_args.hf_hub_token,
            num_proc=data_args.preprocessing_num_workers,
            trust_remote_code=model_args.trust_remote_code,
            streaming=data_args.streaming and dataset_attr.load_from != "file",
        )
        # 如果是从本地文件流式加载，转换为可迭代数据集
        if data_args.streaming and dataset_attr.load_from == "file":
            dataset = dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers)

    # 如果指定了样本数量且非流式模式，随机采样固定数量样本
    if dataset_attr.num_samples is not None and not data_args.streaming:
        target_num = dataset_attr.num_samples
        # 使用 np.random.permutation 对整个数据集索引进行随机打乱；然后取前 target_num 个样本索引；现无重复抽样（如果数据集大小 ≥ target_num）；
        indexes = np.random.permutation(len(dataset))[:target_num]  # all samples should be included
        target_num -= len(indexes)
        # 如果原始数据集比 target_num 小，那么前面只取了所有样本；target_num减去len(indexes)大于0，就是相当于len(indexes)小了，还需要再抽样。
        if target_num > 0:
            # 随机选择 target_num 个样本
            expand_indexes = np.random.choice(len(dataset), target_num)
            indexes = np.concatenate((indexes, expand_indexes), axis=0)

        assert len(indexes) == dataset_attr.num_samples, "Sample num mismatched."
        dataset = dataset.select(indexes)
        logger.info_rank0(f"Sampled {dataset_attr.num_samples} examples from dataset {dataset_attr}.")

    # 如果设置了最大样本数，截断数据集
    if data_args.max_samples is not None:  # truncate dataset
        max_samples = min(data_args.max_samples, len(dataset))
        dataset = dataset.select(range(max_samples))

    # 最后将数据集对齐到统一格式，返回标准结构
    return align_dataset(dataset, dataset_attr, data_args, training_args)


# 合并数据集
# 合并数据集
def _get_merged_dataset(
    dataset_names: Optional[list[str]],
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    stage: Literal["pt", "sft", "rm", "ppo", "kto"],
    return_dict: bool = False,
) -> Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]]:
    r"""Return the merged datasets in the standard format."""
    # 如果未指定任何数据集名称，直接返回 None
    if dataset_names is None:
        return None

    # 初始化一个空字典用于存储每个数据集名称对应的数据集对象
    datasets = {}
    # 遍历每个数据集名称及其对应的属性配置
    for dataset_name, dataset_attr in zip(dataset_names, get_dataset_list(dataset_names, data_args.dataset_dir)):
        # 检查当前训练阶段是否与该数据集是否支持 ranking 标签匹配
        if (stage == "rm" and dataset_attr.ranking is False) or (stage != "rm" and dataset_attr.ranking is True):
            raise ValueError("The dataset is not applicable in the current training stage.")

        # 加载单个数据集，并保存到字典中
        datasets[dataset_name] = _load_single_dataset(dataset_attr, model_args, data_args, training_args)

    # 如果需要返回字典格式，则直接返回数据集字典
    if return_dict:
        return datasets
    else:
        # 否则将所有数据集合并为一个统一的大数据集
        return merge_dataset(list(datasets.values()), data_args, seed=training_args.seed)


def _get_dataset_processor(
    data_args: "DataArguments",
    stage: Literal["pt", "sft", "rm", "ppo", "kto"],
    template: "Template",
    tokenizer: "PreTrainedTokenizer",
    processor: Optional["ProcessorMixin"],
    do_generate: bool = False,
) -> "DatasetProcessor":
    r"""Return the corresponding dataset processor."""
    # 如果训练阶段是预训练 (pt)，使用 PretrainDatasetProcessor
    if stage == "pt":
        dataset_processor_class = PretrainDatasetProcessor
    # 如果是监督微调 (sft) 且不需要生成，则根据 packing 设置选择处理器
    elif stage == "sft" and not do_generate:
        # 如果启用了数据打包（packing），则进一步判断是否使用 neat_packing 特性
        if data_args.packing:
            # 如果启用 neat_packing，临时打补丁优化 attention mask 的类型处理
            if data_args.neat_packing:  # hack datasets to have int32 attention mask
                from datasets.arrow_writer import OptimizedTypedSequence, TypedSequence

                def __init__(self, data, **kwargs):
                    return TypedSequence.__init__(
                        self,
                        data,
                        type=kwargs.pop("type", None),
                        try_type=kwargs.pop("try_type", None),
                        optimized_int_type=kwargs.pop("optimized_int_type", None),
                    )

                OptimizedTypedSequence.__init__ = __init__
            # 使用打包格式的监督学习处理器
            dataset_processor_class = PackedSupervisedDatasetProcessor
        else:
            # 不打包时使用普通监督学习处理器
            dataset_processor_class = SupervisedDatasetProcessor

    # 如果是奖励模型训练 (rm)，使用 PairwiseDatasetProcessor
    elif stage == "rm":
        dataset_processor_class = PairwiseDatasetProcessor
    # 如果是 KTO 算法 (kto)，使用 FeedbackDatasetProcessor
    elif stage == "kto":
        dataset_processor_class = FeedbackDatasetProcessor
    # 默认情况使用无监督处理器（如 PPO 或其他未匹配的 stage）
    else:
        dataset_processor_class = UnsupervisedDatasetProcessor

    # 实例化并返回对应的 DatasetProcessor 子类
    return dataset_processor_class(template=template, tokenizer=tokenizer, processor=processor, data_args=data_args)


def _get_preprocessed_dataset(
    dataset: Optional[Union["Dataset", "IterableDataset"]],
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    stage: Literal["pt", "sft", "rm", "ppo", "kto"],
    template: "Template",
    tokenizer: "PreTrainedTokenizer",
    processor: Optional["ProcessorMixin"] = None,
    is_eval: bool = False,
) -> Optional[Union["Dataset", "IterableDataset"]]:
    r"""Preprocesses the dataset, including format checking and tokenization."""
    # 如果传入数据集为空，直接返回 None
    if dataset is None:
        return None

    # 获取对应训练阶段的数据集处理器，负责格式转换、tokenize 等操作
    dataset_processor = _get_dataset_processor(
        data_args, stage, template, tokenizer, processor, do_generate=(training_args.predict_with_generate and is_eval)
    )
    # 获取数据集中的列名，用于后续删除原始字段
    column_names = list(next(iter(dataset)).keys())
    # 初始化 map 预处理参数字典
    kwargs = {}
    # 非流式加载时，设置多进程、缓存等参数
    if not data_args.streaming:
        kwargs = dict(
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=(not data_args.overwrite_cache) or (training_args.local_process_index != 0),
            desc="Running tokenizer on dataset",
        )

    # 对数据集进行 map 操作，执行预处理函数（tokenize、格式检查等）
    dataset = dataset.map(
        dataset_processor.preprocess_dataset,
        batched=True,
        batch_size=data_args.preprocessing_batch_size,
        remove_columns=column_names,
        **kwargs,
    )

    # 如果当前是日志阶段，则打印一个示例查看预处理结果
    if training_args.should_log:
        try:
            print("eval example:" if is_eval else "training example:")
            dataset_processor.print_data_example(next(iter(dataset)))
        except StopIteration:
            # 如果找不到有效样本，抛出异常提示用户检查数据或扩大数据集
            if stage == "pt":
                raise RuntimeError("Cannot find sufficient samples, consider increasing dataset size.")
            else:
                raise RuntimeError("Cannot find valid samples, check `data/README.md` for the data format.")

    # 返回预处理后的数据集
    return dataset


# 处理数据的主程序
def get_dataset(
    template: "Template",
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    stage: Literal["pt", "sft", "rm", "ppo", "kto"],
    tokenizer: "PreTrainedTokenizer",
    processor: Optional["ProcessorMixin"] = None,
) -> "DatasetModule":
    r"""Get the train dataset and optionally gets the evaluation dataset."""
    # Load tokenized dataset if path exists
    # 如果指定了 tokenized_path 并且该路径下存在数据集，则尝试从磁盘加载
    if data_args.tokenized_path is not None:
        # 判断指定路径是否存在可加载的 tokenized 数据集
        if has_tokenized_data(data_args.tokenized_path):
            # 加载数据集时忽略其他数据参数
            logger.warning_rank0("Loading dataset from disk will ignore other data arguments.")
            # 从磁盘加载 tokenized 数据集
            tokenized_data = load_from_disk(data_args.tokenized_path)
            # 将数据封装为 DatasetModule 格式
            dataset_module = get_dataset_module(tokenized_data)
            # 若启用流式处理，将训练集转为可迭代数据集
            if data_args.streaming:
                dataset_module["train_dataset"] = dataset_module["train_dataset"].to_iterable_dataset()

            # 输出加载成功信息
            logger.info_rank0(f"Loaded tokenized dataset from {data_args.tokenized_path}.")
            # 返回封装后的数据模块
            return dataset_module

        # 如果启用 streaming 模式但要保存到磁盘，抛出错误
        if data_args.streaming:
            raise ValueError("Turn off `streaming` when saving dataset to disk.")

    # Load and preprocess dataset
    # 在主进程中加载原始数据集（多个数据源合并）
    with training_args.main_process_first(desc="load dataset", local=(not data_args.data_shared_file_system)):
        # 获取训练用的合并数据集
        dataset = _get_merged_dataset(data_args.dataset, model_args, data_args, training_args, stage)
        # 获取评估用的合并数据集（可能包含多个子集）
        eval_dataset = _get_merged_dataset(
            data_args.eval_dataset,
            model_args,
            data_args,
            training_args,
            stage,
            return_dict=data_args.eval_on_each_dataset,
        )

    # 在主进程中对数据进行预处理（tokenize、格式转换等）
    with training_args.main_process_first(desc="pre-process dataset", local=(not data_args.data_shared_file_system)):
        # 对训练集进行预处理
        dataset = _get_preprocessed_dataset(
            dataset, data_args, training_args, stage, template, tokenizer, processor, is_eval=False
        )
        # 如果是字典形式的评估集，分别对每个子集进行预处理
        if isinstance(eval_dataset, dict):
            for eval_name, eval_data in eval_dataset.items():
                eval_dataset[eval_name] = _get_preprocessed_dataset(
                    eval_data, data_args, training_args, stage, template, tokenizer, processor, is_eval=True
                )
        else:
            # 否则直接对普通评估集进行预处理
            eval_dataset = _get_preprocessed_dataset(
                eval_dataset, data_args, training_args, stage, template, tokenizer, processor, is_eval=True
            )

        # 分割或组织最终的训练和验证数据集
        dataset_dict = split_dataset(dataset, eval_dataset, data_args, seed=training_args.seed)
        # 如果指定了 tokenized_path，将处理好的数据保存到磁盘
        if data_args.tokenized_path is not None:  # save tokenized dataset to disk
            if training_args.should_save:
                # 保存数据集到磁盘路径
                dataset_dict.save_to_disk(data_args.tokenized_path)
                # 记录保存成功信息
                logger.info_rank0(f"Tokenized dataset is saved at {data_args.tokenized_path}.")
                # 提示用户下次可以直接使用 tokenized_path 进行训练
                logger.info_rank0(f"Please launch the training with `tokenized_path: {data_args.tokenized_path}`.")

        # 返回封装成 DatasetModule 的数据结构
        return get_dataset_module(dataset_dict)
